import torch
from cifa_nn import Net
import cv2
import numpy as np

class Transform(object):
    def __init__(self):
        self.model = Net()

    def Load(self):
        self.model.load_state_dict(torch.load('/home/nvidia/桌面/'
                                         'cifar10颜色识别/cifar2_img32*32_officialtest'))

    def trans(self, img):
        image1 = cv2.resize(img, (32, 32))
        image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
        image1 = np.transpose(image1, (2, 0, 1))
        T_images = torch.Tensor(image1)
        T_images = T_images.unsqueeze(0)
        return T_images

    def predict(self, T_images):
        outputs = self.model(T_images)
        _, predicted = torch.max(outputs, 1)
        return predicted
